# 自主式 SQL 智能体：具备自主探索、分步执行及任务终结核心能力
# 自主式 SQL 智能体：具备自主探索、分步执行及任务终结核心能力
# 简介
# 本文面向技术爱好者与开发者，系统性讲解一种无需人工 Schema、具备自主探索、分步执行与任务终结能力的自主式 SQL 智能体（Autonomous SQL Agent）。文章包含技术原理、专业流程图、完整代码实现

# 核心能力
# 在传统 Text-to-SQL 的工作流中，往往需要人工提供数据库 Schema，且往往需要开发者对 SQL 的执行结果进行人工分析与纠错。自主式 SQL 智能体的目标是实现“零人工 Schema 依赖、逐步自研 SQL、自动错误修复、任务终结明确化”的端到端能力。核心能力包括：

# 自主探索：AI 通过对数据库元数据的查询，自动推断出可操作对象及结构，无需人工提供 Schema。
# 分步执行：将复杂任务拆解成单步 SQL 操作，逐步推进，降低因单步失败带来的整体风险。
# 自动纠错：若某轮 SQL 执行失败，智能体能够分析错误原因并给出修正后的 SQL。
# 任务终结：在任务完成时输出明确的结束标记，便于后续自动化流水线的检测。
# 流式交互：采用对话式的任务执行，过程透明、可控。
# 工作流程要点
# 用户输入任务 → AI 生成首个 SQL 及下一步执行计划（以...包裹）。
# 执行 SQL：如执行成功，返回结果并进入下一步；若失败，进入错误分析阶段并生成修正 SQL。
# 循环迭代：直到任务结束信号出现（）。
# 任务结束输出：最终输出明确的标记，方便结束任务。
# 流程图

# graph TD
#   A[用户输入任务] --> B[AI 生成首个 SQL]
#   B --> C{执行 SQL}
#   C -->|成功| D[返回结果]
#   C -->|失败| E[返回错误]
#   D --> F[AI分析结果]
#   E --> G[AI分析错误原因]
#   F --> H{任务完成？}
#   G --> B
#   H -->|是| I[输出 <END>]
#   H -->|否| B

# 完整代码实现
# 以下代码实现一个完整的自主式 SQL 智能体。核心思想与文末测试案例一致，代码实现遵循以下设计原则：

# 不依赖人工 Schema 自动探索表结构；
# 将任务拆解为单步 SQL，逐步推进；
# SQL 执行具备安全保护，查询型语句允许返回数据，其他操作返回影响行数；
# 遇到错误自动分析并给出修正 SQL；
# 支持流式输出，便于实现对话式交互。
# 请将下列代码保存为 AgentSQL.py，以便本文章中的演示与验证。


# AgentSQL.py
# 版本：自主式 SQL 智能体实现（基于 SQLite，面向教学与研究）
# 说明：本实现参考本文所描述的设计思路，核心能力包含自主探索、分步执行、自动修复与任务结束标识。

import sqlite3
import os
import sys
import re
from typing import Tuple, Dict, List, Generator, Optional

# 第三方依赖：OpenAI
# try:
#     from openai import OpenAI
# except Exception:
#     # 兼容环境：未安装 openai 时的降级处理
#     OpenAI = None  # type: ignore

# 读取环境变量（密钥等配置）
from dotenv import load_dotenv
load_dotenv()

# 数据库与 AI 配置
# DB_NAME = "company.db"  # SQLite  数据库文件名
DB_NAME = "Chinook.db"
# API_KEY = os.getenv("API_KEY", "")
# BASE_URL = os.getenv("BASE_URL", "")
# MODEL_NAME = os.getenv("MODEL_NAME", "gpt-3.5-turbo")
MODEL_NAME = os.getenv("MODEL_NAME", "qwen3:8b")

# OpenAI 客户端（若未配置 API，则在运行时抛错或可用本地模拟）
# if OpenAI is not None and API_KEY:
#     client = OpenAI(api_key=API_KEY, base_url=BASE_URL)  # type: ignore
# else:
#     client = None  # 运行时将提示缺少 OpenAI 客户端

from langchain_ollama import ChatOllama
client = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False)

# 系统提示词：引导 AI 进行自主探索、分步执行与错误修复
SYSTEM_PROMPT = """
你是一个自主式数据库智能体，必须按以下流程完成任务：

### 🧩 工作流程（必须严格遵守）：
1. 探索阶段：若不知道表结构，先用 SQL 查询来获取元数据。
   - 查表名：SELECT name FROM sqlite_master WHERE type='table';
   - 查看表结构：PRAGMA table_info(<table_name>);
2. 分步执行：将任务拆分为多个步骤，每轮只执行一个 SQL。
3. 反馈驱动：每执行完一个 SQL，我会把结果返回给你。
4. 终结标识：当所有任务完成时，在最后一句加上 <END>。
5. 错误修复：如果执行失败，分析错误并生成修正 SQL。请给出修正后的 SQL 并等待我确认后执行。

### 📝 输出格式要求：
- 每个 SQL 用 <sql>...</sql> 包围
- 每次只生成一个 SQL（除非明确要求批量）
- 正常聊天时不要用 SQL 标签
- 任务完成必须加 <END>

### 🚫 禁止行为：
- 不要假设表结构
- 不要一次性生成多个 SQL
- 不要编造字段名

现在开始，等待用户指令。
"""

class AutonomousSQLAgent:
    def __init__(self, db_name: str = DB_NAME):
        self.db_name = db_name
        self.chat_history: List[Dict[str, str]] = [
            {"role": "system", "content": SYSTEM_PROMPT}
        ]

    def stream_generate(self, user_input: str) -> Generator[str, None, None]:
        """
        流式生成 AI 响应（简化版：依赖 OpenAI 流式接口）
        """
        self.chat_history.append({"role": "user", "content": user_input})
        full_content = ""

        if client is None:
            yield "\n警告：未配置 OpenAI 客户端，无法实际对话，请在环境中正确配置 API_KEY、MODEL_NAME 等。"
            return

        try:
            print("开始生成响应...")
            # print(self.chat_history)
            response = client.stream(input=self.chat_history)
            # response = client.chat.completions.create(  # type: ignore[attr-defined]
            #     model=MODEL_NAME,
            #     messages=self.chat_history,
            #     stream=True,
            # )

            for chunk in response:
                # print(chunk)
                # 兼容不同返回结构
                content = ""
                # 1. 针对新版 OpenAI SDK
                if isinstance(chunk, dict):
                    # print("chunk is dict")
                    content = chunk.content
                    # choices = chunk.get("choices", [])
                    # if choices:
                    #     delta = choices[0].get("delta", {})
                    #     if isinstance(delta, dict):
                    #         content = delta.get("content", "")
                else:
                    # 对象属性访问
                    # print("chunk is not dict")
                    try:
                        # content = chunk.choices[0].delta.content  # type: ignore
                        content = chunk.content  # type: ignore
                    except Exception:
                        content = ""

                if content:
                    full_content += content
                    yield content

            # print(full_content)
            self.chat_history.append({"role": "assistant", "content": full_content})

        except Exception as e:
            error_msg = f"\n❌ API 错误：{str(e)}"
            yield error_msg
            self.chat_history.append({"role": "assistant", "content": error_msg})

    def extract_sql(self, text: str) -> str:
        """从文本中提取 <sql>...</sql> 内的语句"""
        m = re.search(r"<sql>(.*?)</sql>", text, re.DOTALL | re.IGNORECASE)
        return m.group(1).strip() if m else ""

def _format_result(headers: List[str], rows: List[Tuple]) -> str:
    if not rows:
        return "无匹配数据。"
    table = "| " + " | ".join(headers) + " |\n"
    table += "|" + "|".join(["---"] * len(headers)) + "|\n"
    for row in rows:
        table += "| " + " | ".join(str(v) for v in row) + " |\n"
    return table

def execute_sql(sql: str, db_name: str = DB_NAME) -> Tuple[bool, str]:
    """
    安全执行 SQL 并返回结果：
    - SELECT/PRAGMA/EXPLAIN/SHOW 等返回数据表格
    - 其他写操作返回受影响行数
    """
    try:
        conn = sqlite3.connect(db_name)
        cursor = conn.cursor()
        print(f"执行 SQL -> {sql}")
        cursor.execute(sql)

        lower = sql.strip().lower()
        if lower.startswith(("select", "pragma", "explain", "show")):
            rows = cursor.fetchall()
            if not rows:
                result = "无匹配数据。"
            else:
                headers = [desc[0] for desc in cursor.description]  # type: ignore
                result = _format_result(headers, rows)
        else:
            changes = conn.total_changes
            result = f"执行成功，影响 {changes} 行。"

        conn.commit()
        conn.close()
        return True, result
    except Exception as e:
        return False, f"执行失败：{str(e)}"

def _build_exploration_sql(table_name: str) -> str:
    return f"PRAGMA table_info({table_name});"

def _list_tables_sql() -> str:
    return "SELECT name FROM sqlite_master WHERE type='table';"

def _ensure_table_names_exist(sql_text: str) -> List[str]:
    """
    简单工具：解析当前数据库中的表名，用于辅助解释阶段。
    实际环境中可增强为对 schema 的完整遍历。
    """
    try:
        conn = sqlite3.connect(DB_NAME)
        cursor = conn.cursor()
        cursor.execute(_list_tables_sql())
        rows = cursor.fetchall()
        conn.close()
        return [row[0] for row in rows] if rows else []
    except Exception:
        return []

class AutonomousSQLAgentRunner:
    """
    提供一个简化的 CLI 入口，用于演示 AutonomousSQLAgent 的工作流程。
    """
    def __init__(self, db_name: str = DB_NAME):
        self.agent = AutonomousSQLAgent(db_name=db_name)

    def run(self):
        print("🚀 欢迎使用自主式 SQL 智能体（支持自主探索 + 分步执行 + 任务终结）")
        print("💡 示例：查询技术部薪资最高的员工")
        print("输入 quit/exit/退出 结束。")
        print("-" * 60)

        while True:
            try:
                user_input = input("\n💬 你: ").strip()
                if user_input.lower() in ["quit", "exit", "退出"]:
                    print("👋 再见！")
                    break
                if not user_input:
                    continue

                # 任务：生成首轮 SQL
                print("🤖 助手: ", end="", flush=True)
                full_response = ""
                for token in self.agent.stream_generate(user_input):
                    print(token, end="", flush=True)
                    full_response += token
                    if "<sql>" in full_response and "</sql>" in full_response:
                        # 已生成可执行的 SQL，进入执行阶段
                        break
                
                # print("SQL", full_response)

                if "<sql>" not in full_response:
                    print()
                    continue

                # 循环执行直到任务结束
                while True:
                    sql = self.agent.extract_sql(full_response)
                    if not sql:
                        break

                    print(f"\n\n⚙️ 正在执行 SQL...", full_response)
                    success, result = execute_sql(sql)
                    if success:
                        print(f"✅ 执行成功！\n🔍 结果：\n{result}")
                        feedback = f"""
SQL 执行成功。
执行的 SQL：
{sql}

执行结果：
{result}

请继续下一步，或如果任务已完成，请回复并加上 <END>。
"""
                    else:
                        print(f"❌ 执行失败：{result}")
                        feedback = f"""
SQL 执行失败。
执行的 SQL：
{sql}

错误信息：
{result}

请分析错误并生成修正后的 SQL。
"""

                    # 下一步：将反馈再发送给 AI 进行下一轮推导
                    print("\n🤖 助手: ", end="", flush=True)

                    full_response = ""
                    for token in self.agent.stream_generate(feedback):
                        print(token, end="", flush=True)
                        full_response += token
                        if "<END>" in full_response:
                            print("\n🎯 任务已完成。")
                            break
                        if "<sql>" in full_response and "</sql>" in full_response:
                            break

                    if "<END>" in full_response:
                        break
                    # 防止极端循环
                    if len(full_response) > 100000:
                        print("\n⚠️  长文本过长，中止任务。")
                        break

            except KeyboardInterrupt:
                print("\n\n👋 再见！")
                break
            except Exception as e:
                print(f"\n\n❌ 程序错误：{str(e)}")

# 仅用于直接运行测试的入口
if __name__ == "__main__":
    # 需要用户自行在项目根目录放置 company.db 数据库
    runner = AutonomousSQLAgentRunner(db_name=DB_NAME)
    runner.run()

# 上述实现了一个简化的自主式 SQL 智能体框架，核心组件与前文描述一致，包含：
# 自主探索：通过查询 sqlite_master 获取表名，使用 PRAGMA table_info 获取字段信息，避免人工提供 Schema。
# 分步执行：任务以单轮 SQL 形式推进，每轮只执行一个 SQL 语句。
# 自动修复：当 SQL 执行失败时，自动分析错误信息并给出修正后的 SQL（在实际对话中以字符串形式提供给 AI，等待下一轮迭代）。
# 任务终结：通过检测输出中包含来判定任务结束。
# 流式输出：通过 stream_generate 实现对话式、逐步输出，提升交互体验。

# | Album |
# | Artist |
# | Customer |
# | Employee |
# | Genre |
# | Invoice |
# | InvoiceLine |
# | MediaType |
# | Playlist |
# | PlaylistTrack |
# | Track |

# - CustomerId (INTEGER, 不可为空, 主键)
# - FirstName (NVARCHAR(40), 不可为空)
# - LastName (NVARCHAR(20), 不可为空)
# - Company (NVARCHAR(80))
# - Address (NVARCHAR(70))
# - City (NVARCHAR(40))
# - State (NVARCHAR(40))
# - Country (NVARCHAR(40))
# - PostalCode (NVARCHAR(10))
# - Phone (NVARCHAR(24))
# - Fax (NVARCHAR(24))
# - Email (NVARCHAR(60), 不可为空)
# - SupportRepId (INTEGER)